/*
 * The MIT License
 * 
 * Copyright (c) 2008 Comcast Interactive Media
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package cim.sf.core.command.sh;

import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.ArrayList;
import java.util.List;

import cim.sf.core.command.CommandInvocation;
import cim.sf.core.command.CommandNotFoundException;
import cim.sf.core.command.CommandShell;

/**
 * A simple Command-executing protocol implementation that works like this:
 * 
 * Client (client.pl) connects and sends the following to start command execution:
 * <ol>
 * <li>The string "SCMD"</li>
 * <li>The contents of the ".key" file located alongside the client script followed by '\0'</li>
 * <li>The client's current working directory followed by '\0'</li>
 * <li>The command path (e.g., /equinox/ss) followed by '\0'</li>
 * <li>For each argument on the script's command line, the argument followed by '\0'</li>
 * <li>'\0'</li>
 * </ol>
 * 
 * At this point, the client will start sending anything passed to STDIN if the "-" parameter
 * is present on its command line. Otherwise, it will shutdown its output stream.
 * <p>
 * The command can respond to the client in three ways:
 * <ul>
 * <li>Output stream (STDOUT)</li>
 * <li>Error stream (STDERR)</li>
 * <li>Status/return code</li>
 * </ul>
 * The server combines these three streams into one to send back to the client. When the command writes
 * bytes to one of these streams, the bytes are prefixed with the following and sent down the 
 * socket output stream to the client:
 * <ol>
 * <li>A 1-byte prefix ('O' for output, 'E' for error, 'S' for status/return code)
 * <li>A 4-byte length of the bytes following
 * </ol>
 * 
 * The server will close the command socket when the {@link cim.sf.core.command.Command#execute(CommandInvocation)}
 * method returns.
 * 
 * Before the ".key" file is sent, the server uses a read timeout of 5 seconds. After the key is sent,
 * this timeout is bumped up to 1 minute to allow for slow STDIN streams on the client side.
 * 
 * @author <a href="mailto:jeremy_volkman@cable.comcast.com">Jeremy Volkman</a>
 */
public class CommandServer {
    private Activator activator;
    private ServerSocket listener;
    private Thread acceptorThread;
    
    public CommandServer(Activator activator) {
        this.activator = activator;
    }
    
    public void open() {
        try {
            listener = new ServerSocket(activator.port, 50, InetAddress.getByName(null));
            acceptorThread = new Thread(new Acceptor(listener), "SHCMD Acceptor");
            acceptorThread.setDaemon(true);
            acceptorThread.start();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public void close() {
        if (listener != null) {
            acceptorThread.interrupt();
            acceptorThread = null;
            try {
                listener.close();
            } catch (IOException e) {
            }
            listener = null;
        }
    }

    private class Acceptor implements Runnable {
        
        private ServerSocket socket;
        
        private Acceptor(ServerSocket socket) {
            this.socket = socket;
        }
        
        public void run() {
            while (!Thread.currentThread().isInterrupted()) {
                try {
                    Socket connection = socket.accept();
                    Executor e = new Executor(connection);
                    Thread t = new Thread(e, "SHCMD");
                    t.setDaemon(false);
                    t.start();
                } catch (IOException e) {
                    try {
                        // don't loop continuously if something goes wrong
                        Thread.sleep(500);
                    } catch (InterruptedException f) {
                        break;
                    }
                }
            }
        }
    }

    private class Executor implements Runnable {

        private static final int expectedIdent = 0x53434D44; // "SCMD"
        
        private Socket socket;

        private Executor(Socket socket) {
            this.socket = socket;
        }

        public void run() {
            ByteArrayOutputStream buf = new ByteArrayOutputStream();

            String cwd = null;
            String commandPath = null;
            
            List<String> params = new ArrayList<String>();
            boolean parseSuccessful = false;
            boolean stdinPresent = false;
            try {
                socket.setSoTimeout(5000); // 5 second read timeout pre-key
                InputStream is = socket.getInputStream();
                DataInputStream dis = new DataInputStream(is);
                
                int ident = dis.readInt();
                if (ident != expectedIdent) {
                    // Not the expected client.
                    socket.close();
                    return;
                }
                
                // first read key.
                // read a max of activator.key.length() bytes.
                
                byte[] keybuf = new byte[activator.key.length() + 1];
                int rlen = 0;
                while (rlen < keybuf.length) {
                    int n = is.read(keybuf, rlen, keybuf.length - rlen);
                    if (n == -1) break;
                    rlen += n;
                }
                
                String key = new String(keybuf, 0, keybuf.length - 1);
                if (!key.equals(activator.key)) {
                    // No way.
                    socket.close();
                    return;
                }

                socket.setSoTimeout(60000); // bump read timeout up to 1 minute
                
                int c;
                while ((c = is.read()) != -1) {
                    if (c == 0) {
                        if (buf.size() > 0) {
                            // End of token
                            if (cwd == null) {
                                cwd = buf.toString();
                            } else if (commandPath == null) {
                                commandPath = buf.toString();
                            } else {
                                String param = buf.toString();
                                if ("-".equals(param)) {
                                    stdinPresent = true;
                                } else {
                                    params.add(buf.toString());
                                }
                            }
                            buf.reset();
                        } else {
                            parseSuccessful = commandPath != null;
                            break;
                        }
                    } else {
                        buf.write(c);
                    }
                }
                if (parseSuccessful) {
                    PrintStream out = new PrintStream(new PrefixedOutputStream(socket.getOutputStream(), 'O'), true);
                    PrintStream err = new PrintStream(new PrefixedOutputStream(socket.getOutputStream(), 'E'), true);
                    OutputStream status = new PrefixedOutputStream(socket.getOutputStream(), 'S');
                    CommandInvocation ci = new CommandInvocation(out, err, stdinPresent ? is : null, new File(cwd), params.toArray(new String[params.size()]));
                    Thread.currentThread().setName("SHCMD - " + commandPath);

                    CommandShell shell = activator.getCommandShell();
                    int sc = 1;
                    if (shell == null) {
                        ci.err.println("CommandShell unavailable.");
                    } else {
                        try {
                            sc = shell.executeCommand(commandPath, ci);
                        } catch (CommandNotFoundException e) {
                            err.println("Command not found: " + commandPath);
                        }
                    }
                    status.write(String.valueOf(sc).getBytes());
                    status.flush();
                }
            } catch (IOException e) {
                // Don't do anything, we'll just close the connection in the finally.
            } finally {
                try {
                    socket.close();
                } catch (IOException e) {
                }
            }
        }
    }

    private static class PrefixedOutputStream extends OutputStream {
        
        private char prefix;
        private OutputStream os;
        private DataOutputStream dos;
        
        public PrefixedOutputStream(OutputStream os, char prefix) {
            this.os = os;
            this.dos = new DataOutputStream(os);
            this.prefix = prefix;
        }
        
        
        @Override
        public void write(int b) throws IOException {
            synchronized(os) {
                dos.writeByte(prefix);
                dos.writeInt(1);
                dos.writeByte(b);
            }
        }
        
        public void write(byte b[], int off, int len) throws IOException {
            if (b == null) {
                throw new NullPointerException();
            } else if ((off < 0) || (off > b.length) || (len < 0) ||
                   ((off + len) > b.length) || ((off + len) < 0)) {
                throw new IndexOutOfBoundsException();
            } else if (len == 0) {
                return;
            }
            synchronized(os) {
                dos.writeByte(prefix);
                dos.writeInt(len);
                dos.write(b, off, len);
            }
        }
        
        public void close() throws IOException {
            dos.close();
        }
        
        public void flush() throws IOException {
            synchronized(os) {
                dos.flush();
            }
        }
    }
}


